import numpy as np
from config import *


class MazeEnvironment:
    def __init__(self, size=ENV_SIZE, inner_size=INNER_SIZE, num_coins=NUM_COINS,
                 num_obstacles=NUM_OBSTACLES, seed=RANDOM_SEED):
        """初始化环境"""
        self.size = size
        self.inner_size = inner_size
        self.num_coins = num_coins
        self.num_obstacles = num_obstacles
        self.seed = seed

        # 初始化随机数生成器
        self.rng = np.random.RandomState(seed)

        # 初始化网格和智能体位置
        self.grid = np.zeros((size, size), dtype=int)
        self.agent_pos = None
        self.start_pos = None
        self.end_pos = None
        self.coins = []

        # 初始化环境
        self.initialize_environment()

    def initialize_environment(self):
        """初始化环境并放置物体"""
        # 用边界填充网格
        self.grid.fill(EMPTY)

        # 计算内部网格的起始位置（居中）
        start_row = (self.size - self.inner_size) // 2
        start_col = (self.size - self.inner_size) // 2

        # 在内部网格的四个角落添加起点、终点和金币
        corners = [
            (start_row, start_col),  # 左上
            (start_row, start_col + self.inner_size - 1),  # 右上
            (start_row + self.inner_size - 1, start_col),  # 左下
            (start_row + self.inner_size - 1, start_col + self.inner_size - 1)  # 右下
        ]

        # 随机打乱角落顺序
        self.rng.shuffle(corners)

        # 放置起点、终点
        self.start_pos = corners[0]
        self.end_pos = corners[1]
        self.grid[self.start_pos] = START
        self.grid[self.end_pos] = END

        # 初始时，agent在起点
        self.agent_pos = self.start_pos
        self.grid[self.agent_pos] = AGENT

        # 放置金币
        self.coins = []
        valid_positions = []
        for i in range(start_row, start_row + self.inner_size):
            for j in range(start_col, start_col + self.inner_size):
                if self.grid[i, j] == EMPTY:
                    valid_positions.append((i, j))

        # 确保有足够的有效位置
        num_coins = min(self.num_coins, len(valid_positions))
        if num_coins > 0:
            coin_positions = self.rng.choice(len(valid_positions), num_coins, replace=False)
            for pos_idx in coin_positions:
                pos = valid_positions[pos_idx]
                self.grid[pos] = COIN
                self.coins.append(pos)
                valid_positions.remove(pos)

        # 放置障碍物
        num_obstacles = min(self.num_obstacles, len(valid_positions))
        if num_obstacles > 0:
            obstacle_positions = self.rng.choice(len(valid_positions), num_obstacles, replace=False)
            for pos_idx in obstacle_positions:
                pos = valid_positions[pos_idx]
                self.grid[pos] = OBSTACLE

    def reset(self):
        """重置环境，将agent放回起点"""
        # 找到当前agent位置并清除
        if self.agent_pos is not None:
            self.grid[self.agent_pos] = EMPTY

        # 将agent放回起点
        self.agent_pos = self.start_pos
        self.grid[self.agent_pos] = AGENT

        return self.get_state()

    def reset_without_changing_layout(self):
        """重置agent位置但保持环境布局不变"""
        # 找到当前agent位置并清除
        if self.agent_pos is not None:
            self.grid[self.agent_pos] = EMPTY

        # 将agent放回起点
        self.agent_pos = self.start_pos
        self.grid[self.agent_pos] = AGENT

        return self.get_state()

    def get_state(self):
        """返回当前环境状态"""
        return self.grid.copy()

    def is_valid_move(self, pos):
        """检查移动是否有效"""
        # 检查是否在网格范围内
        if not (0 <= pos[0] < self.size and 0 <= pos[1] < self.size):
            return False

        # 检查是否是障碍物
        return self.grid[pos] != OBSTACLE

    def get_next_state_after_action(self, action):
        """获取执行动作后的下一个状态，但不实际改变环境"""
        # 创建网格副本
        next_grid = self.grid.copy()

        # 存储当前位置
        old_pos = self.agent_pos

        # 根据动作确定新位置
        new_pos = list(old_pos)
        if action == UP:
            new_pos[0] -= 1
        elif action == DOWN:
            new_pos[0] += 1
        elif action == LEFT:
            new_pos[1] -= 1
        elif action == RIGHT:
            new_pos[1] += 1
        elif action == STAY:
            pass

        new_pos = tuple(new_pos)

        # 检查新位置是否有效
        if self.is_valid_move(new_pos):
            # 更新网格
            next_grid[old_pos] = EMPTY
            next_grid[new_pos] = AGENT

        return next_grid

    def step(self, action):
        """执行动作并返回更新后的状态"""
        # 存储当前位置
        old_pos = self.agent_pos

        # 根据动作确定新位置
        new_pos = list(old_pos)
        if action == UP:
            new_pos[0] -= 1
        elif action == DOWN:
            new_pos[0] += 1
        elif action == LEFT:
            new_pos[1] -= 1
        elif action == RIGHT:
            new_pos[1] += 1
        elif action == STAY:
            pass

        new_pos = tuple(new_pos)

        # 检查新位置是否有效
        if self.is_valid_move(new_pos):
            # 创建网格副本
            grid_copy = self.grid.copy()

            # 如果新位置有金币，收集它
            if new_pos in self.coins:
                self.coins.remove(new_pos)

            # 更新网格
            grid_copy[old_pos] = EMPTY
            grid_copy[new_pos] = AGENT

            # 更新状态
            self.grid = grid_copy
            self.agent_pos = new_pos

        return self.get_state()

    def state_to_input(self, state):
        """将状态转换为模型输入格式"""
        # 进行one-hot编码
        encoded = np.zeros((self.size, self.size, CELL_FEATURES))
        for i in range(self.size):
            for j in range(self.size):
                encoded[i, j, state[i, j]] = 1

        # 将编码后的状态展平为一维数组
        return encoded.flatten()

    def state_to_grid(self, state_input):
        """将模型输入格式转换回网格形式"""
        # 对于one-hot编码的状态
        if len(state_input) == self.size * self.size * CELL_FEATURES:
            # 重塑为[height, width, features]
            encoded = state_input.reshape(self.size, self.size, CELL_FEATURES)
            # 转换回网格
            grid = np.argmax(encoded, axis=-1)
            return grid
        else:
            # 尝试将扁平的输入重塑为网格
            try:
                return state_input.reshape(self.size, self.size)
            except:
                # 如果都不成功，返回空网格
                return np.zeros((self.size, self.size), dtype=int)

    def get_agent_position(self, state=None):
        """从状态中获取agent位置"""
        if state is None:
            state = self.grid

        agent_pos = np.where(state == AGENT)
        if len(agent_pos[0]) > 0:
            return (agent_pos[0][0], agent_pos[1][0])
        return None

    def one_hot_encode_state(self, state):
        """将状态one-hot编码"""
        # 创建one-hot编码
        encoded = np.zeros((self.size, self.size, CELL_FEATURES))
        for i in range(self.size):
            for j in range(self.size):
                encoded[i, j, state[i, j]] = 1

        return encoded.flatten()
